import os
import wx
import cv2
import time
import numpy as np
import threading
from concurrent.futures import ThreadPoolExecutor
from core.adb_interface import AdbInterface


class DlgPicture(wx.Dialog):

    class WindowSizeInitType:
        AUTO = 0
        SIZE = 1
        SCALE = 2

    def __init__(self, parent, title="Picture", cv_img: np.ndarray = None, scale: float = None, size: tuple = None):
        wx.Dialog.__init__(self, parent, title=title)
        self.cv_img_original = new_image = cv_img
        self.window_scale = 1.0
        self.image_scale = 1.0
        self.window_size = (800, 600)
        self.flag_init_win_size = self.WindowSizeInitType.AUTO
        # current display
        if parent is None:
            self.monitor = 0
        else:
            self.monitor = wx.Display.GetFromWindow(parent)
            if self.monitor is -1:
                parent_rect: wx.Rect = self.GetScreenRect()
                self.monitor = wx.Display.GetFromPoint((int(parent_rect.GetX() + parent_rect.GetWidth() / 2),
                                                        int(parent_rect.GetY() + parent_rect.GetHeight() / 2)))
                self.monitor = 0 if self.monitor is -1 else self.monitor
        # screen size
        self.desktop_client_rect: wx.Rect = wx.Display(self.monitor).GetClientArea().Get()  # Rect(x,y,w,h)
        self.SetSize(self.desktop_client_rect[2:])
        self.desktop_client_size = self.GetClientSize().Get()
        self.client_border_size = (self.desktop_client_rect[2]-self.desktop_client_size[0],
                                   self.desktop_client_rect[3]-self.desktop_client_size[1])
        # better window size
        if size is not None:
            size = int(size[0]), int(size[1])
            self.flag_init_win_size = self.WindowSizeInitType.SIZE
            if cv_img is not None:
                height, width = cv_img.shape[:2]
                scale = size[0] / width
            else:
                scale = 1.0
        elif scale is not None:
            scale = float(scale)
            self.flag_init_win_size = self.WindowSizeInitType.SCALE
            if cv_img is not None:
                height, width = cv_img.shape[:2]
                size = width * scale, height * scale
            else:
                size = self.desktop_client_size
        else:
            self.flag_init_win_size = self.WindowSizeInitType.AUTO
            if cv_img is not None:
                height, width = cv_img.shape[:2]
                if self.desktop_client_size[0] < width or self.desktop_client_size[1] < height:
                    new_width = min(self.desktop_client_size[0], width)
                    new_height = min(self.desktop_client_size[1], height)
                    size = (int(min(new_height * width / height, new_width)),
                            int(min(new_width * height / width, new_height)))
                    new_image = cv2.resize(cv_img, size)
                    scale = size[0] / width
                else:
                    scale = 1.0
                    size = (width, height)
            else:
                scale = 1.0
                size = self.desktop_client_size
        self.window_size = size
        self.image_scale = 1. / scale
        self.window_scale = scale
        self.cv_img_draw: np.ndarray = new_image
        self.bmp_draw = wx.Bitmap()
        self.draw_image(new_image)
        self.SetClientSize(self.window_size)
        self.SetAutoLayout(True)
        self.Move(self.desktop_client_rect[0] + self.desktop_client_rect[2] / 2 - self.GetSize().GetWidth() / 2,
                  self.desktop_client_rect[1] + self.desktop_client_rect[3] / 2 - self.GetSize().GetHeight() / 2)
        self.Bind(wx.EVT_PAINT, self.on_paint)

    def on_paint(self, e: wx.PaintEvent):
        if self.bmp_draw.IsOk():
            wx.PaintDC(self).DrawBitmap(self.bmp_draw, 0, 0)
        pass

    def reset_client_size(self, size: tuple):
        if self.cv_img_original is not None:
            self.window_scale = size[0] / self.cv_img_original.shape[1]
            self.image_scale = 1. / self.window_scale
        else:
            self.window_scale = 1.
            self.image_scale = 1.
        old_win_size = self.window_size
        old_win_center = list(self.GetPosition().Get())
        old_win_center[0] += int(old_win_size[0] / 2)
        old_win_center[1] += int(old_win_size[1] / 2)
        self.monitor = wx.Display.GetFromWindow(self)
        new_win_pos = [int(old_win_center[0]-size[0]/2),int(old_win_center[1]-size[1]/2)]
        if self.monitor is -1:
            new_win_pos = (max(new_win_pos[0], 0), max(new_win_pos[1], 0))
        else:
            # screen size
            self.desktop_client_rect: wx.Rect = wx.Display(self.monitor).GetClientArea()  # Rect(x,y,w,h)
            self.desktop_client_size = (self.desktop_client_rect.GetWidth()-self.client_border_size[0],
                                        self.desktop_client_rect.GetHeight()-self.client_border_size[1])
            if new_win_pos[0] < self.desktop_client_rect.GetX():
                new_win_pos[0] = self.desktop_client_rect.GetX()
            if new_win_pos[1] < self.desktop_client_rect.GetY():
                new_win_pos[1] = self.desktop_client_rect.GetY()
            if new_win_pos[0] + size[0] > self.desktop_client_rect.GetX() + self.desktop_client_rect.GetWidth():
                new_win_pos[0] = self.desktop_client_rect.GetX()+self.desktop_client_rect.GetWidth()-size[0]
            if new_win_pos[1] + size[1] > self.desktop_client_rect.GetY() + self.desktop_client_rect.GetHeight():
                new_win_pos[1] = self.desktop_client_rect.GetY()+self.desktop_client_rect.GetHeight()-size[1]
        self.window_size = size
        self.SetClientSize(size)
        self.Move(new_win_pos)
        pass

    def set_image(self, cv_img: np.ndarray):
        if cv_img is not None:
            self.window_size = self.GetClientSize().GetWidth(), self.GetClientSize().GetHeight()
            self.cv_img_original = cv_img
            self.window_scale = self.window_size[0] / cv_img.shape[1]
            self.image_scale = 1. / self.window_scale
            self.cv_img_draw = cv2.resize(cv_img, self.window_size)
        else:
            self.cv_img_draw = self.cv_img_original = None
        self.draw_image(self.cv_img_draw)
        pass

    def get_image(self):
        return self.cv_img_original

    def draw_image(self, cv_img: np.ndarray):
        if cv_img is not None:
            h, w = cv_img.shape[:2]
            img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)
            self.bmp_draw = wx.Bitmap.FromBuffer(w, h, img)
            dc = wx.BufferedDC(wx.ClientDC(self), self.bmp_draw)
        else:
            wx.ClientDC(self).Clear()
        pass


class RecordActionForDlgPicture():
    def __init__(self):
        if not isinstance(self, DlgPicture):
            raise RuntimeError("Must be companion class of \'DlgPicture\'.")
        self.clicked_point = [None, None]
        self.device_point = [None, None]
        self.press_time = 0.
        self.press_start_time = 0.
        self.press_stop_time = 0.

    def repaint_mark(self, pos: [list, tuple]):
        if self.cv_img_draw is not None:
            new_pic = self.cv_img_draw.copy()
        else:
            color: wx.Colour = self.GetBackgroundColour()
            new_pic = np.zeros([self.window_size[1], self.window_size[0], 3], np.uint8)
            cv2.rectangle(new_pic, (0, 0), self.window_size, (color.Blue(), color.Green(), color.Red()), -1)
            if self.cv_img_draw is None: self.cv_img_draw = new_pic.copy()
        h, w = new_pic.shape[:2]
        title = "({}, {})".format(int(pos[0] * self.image_scale), int(pos[1] * self.image_scale))
        cv2.line(new_pic, (0, pos[1]), (w, pos[1]), (255, 0, 0), 1)
        cv2.line(new_pic, (pos[0], 0), (pos[0], h), (255, 0, 0), 1)
        if self.clicked_point[0] is not None:
            title += "[({0[0]},{0[1]})".format(self.device_point[0])
            cv2.circle(new_pic, self.clicked_point[0], 4, (0, 255, 0), -1)
            if self.clicked_point[1] is None:
                title += ",{}s".format(time.time() - self.press_start_time)
                cv2.line(new_pic, self.clicked_point[0], pos, (0, 200, 200), 2)
            elif self.clicked_point[1] == self.clicked_point[0]:
                title += ",{}s".format(self.press_time)
                cv2.circle(new_pic, self.clicked_point[0], 10, (0, 0, 255), 2)
            else:
                title += "=>({0[0]},{0[1]}),{1}s".format(self.device_point[1], self.press_time)
                cv2.circle(new_pic, self.clicked_point[1], 4, (0, 0, 255), -1)
                cv2.line(new_pic, self.clicked_point[0], self.clicked_point[1], (0, 220, 220), 2)
            title += "]"
        self.SetTitle(title)
        self.draw_image(new_pic)
        return True

    def on_mouse_up(self, e: wx.MouseEvent):
        if self.clicked_point[0] is not None:
            self.clicked_point[1] = (e.GetX(), e.GetY())
            self.device_point[1] = (int(e.GetX() * self.image_scale), int(e.GetY() * self.image_scale))
        else:
            return False
        if self.press_start_time > 0:
            self.press_stop_time = time.time()
            self.press_time = self.press_stop_time - self.press_start_time
        else:
            return False
        self.repaint_mark((e.GetX(), e.GetY()))
        return True

    def on_mouse_down(self, e: wx.MouseEvent):
        self.clicked_point[0] = e.GetX(), e.GetY()  # not e.GetPosition()->wx.Point
        self.device_point[0] = int(e.GetX() * self.image_scale), int(e.GetY() * self.image_scale)
        self.clicked_point[1] = None
        self.device_point[1] = None
        self.press_start_time = time.time()
        self.press_time = 0.
        self.press_stop_time = 0.
        self.repaint_mark((e.GetX(), e.GetY()))
        return True

    def on_move(self, e: wx.MouseEvent):
        self.repaint_mark((e.GetX(), e.GetY()))
        return True


class DlgScreen(DlgPicture):

    def __init__(self, parent, adb: AdbInterface, scale: float = None, size: tuple = None):
        """
        :param parent:
        :param adb:
        :param scale:
        """
        DlgPicture.__init__(self, parent, title="Screen", cv_img=None, size=None, scale=None)
        # control
        self.lbl_message = wx.StaticText(parent=self, label="正在启动...")
        # parameters
        self.adb = adb
        self.image_size_init = (0,0)
        if size is not None:
            self.reset_client_size(size)
            self.flag_init_win_size = self.WindowSizeInitType.SIZE
            succ, (scr_width, src_height) = self.adb.get_screen_size()
            if succ:
                self.image_size_init = (scr_width, src_height)
            else:
                self.lbl_message.SetLabel("获取设备信息失败！")
        else:
            # cal scale
            succ, (scr_width, src_height) = self.adb.get_screen_size()
            if succ:
                if scale is not None:
                    self.reset_client_size((scr_width * scale, src_height * scale))
                    self.flag_init_win_size = self.WindowSizeInitType.SCALE
                else:
                    self.SetSize(self.desktop_client_size)
                    client_size = self.GetClientSize()
                    win_size = (scr_width, src_height)
                    if client_size[0] < scr_width or client_size[1] < src_height:
                        new_width = min(self.desktop_client_size[0], scr_width)
                        new_height = min(self.desktop_client_size[1], src_height)
                        win_size = (int(min(new_height * scr_width / src_height, new_width)),
                                    int(min(new_width * src_height / scr_width, new_height)))
                    self.reset_client_size(win_size)
                    self.flag_init_win_size = self.WindowSizeInitType.AUTO
            else:
                self.lbl_message.SetLabel("获取设备信息失败！")
                if scale is not None:
                    self.window_scale = scale
                    self.flag_init_win_size = self.WindowSizeInitType.SCALE
                else:
                    self.flag_init_win_size = self.WindowSizeInitType.AUTO
        self.press_time = 0.
        self.press_start_time = 0.

        # layout
        self.box_message = wx.BoxSizer()
        self.box_message.Add(self.lbl_message, proportion=1, flag=wx.ALIGN_CENTER)
        self.box_all = wx.BoxSizer(wx.VERTICAL)
        self.box_all.Add(self.box_message, proportion=1, flag=wx.ALIGN_CENTER)
        self.SetSizer(self.box_all)
        self.SetAutoLayout(True)
        msg_font = wx.Font()
        msg_font_w = max(self.window_size[0] / 20, 15)
        msg_font_h = max(self.window_size[1] / 20, 25)
        msg_font.SetPixelSize((min(msg_font_w, msg_font_h * 0.6), min(msg_font_h, msg_font_w / 0.6)))
        self.lbl_message.SetFont(msg_font)
        # capture
        self.flag_capture = threading.Event()
        self.flag_capture.set()
        self.lock_draw = threading.Lock()
        self.time_draw = time.time()
        self.pool_capture = None
        if self.adb is not None:
            threading.Thread(target=self.init_screen).start()
        self.Bind(wx.EVT_CLOSE, self.on_close)

    def on_close(self, e: wx.CloseEvent):
        if isinstance(self.pool_capture, ThreadPoolExecutor):
            self.flag_capture.clear()
            self.lbl_message.SetLabel("正在关闭...")
            self.lbl_message.Show()
            threading.Thread(target=self.stop_thread_capture).start()

    def capture_and_show(self):
        tm = time.time()
        succ, img = self.adb.capture_screen(img_lib="cv2")
        if succ:
            with self.lock_draw:
                # 丢帧
                if tm < self.time_draw:
                    return
                self.time_draw = tm
                try:
                    self.set_image(img)
                except Exception as ex:
                    print(ex)
                img_size = img.shape[:2]
                # image size changed
                if self.image_size_init[0] != img_size[1] or self.image_size_init[1] != img_size[0]:
                    if self.flag_init_win_size == self.WindowSizeInitType.SCALE:
                        self.reset_client_size((img_size[1]*self.window_scale, img_size[0]*self.window_scale))
                    elif self.flag_init_win_size == self.WindowSizeInitType.SIZE:
                        self.reset_client_size(self.window_size)
                    else:
                        win_size = (img_size[1], img_size[0])
                        if self.desktop_client_size[0] < img_size[1] or self.desktop_client_size[1] < img_size[0]:
                            new_width = min(self.desktop_client_size[0], img_size[1])
                            new_height = min(self.desktop_client_size[1], img_size[0])
                            win_size = (int(min(new_height * img_size[1] / img_size[0], new_width)),
                                        int(min(new_width * img_size[0] / img_size[1], new_height)))
                        self.reset_client_size(win_size)
                    self.image_size_init = (img_size[1], img_size[0])

    def init_screen(self):
        tm_st = time.time()
        self.capture_and_show()
        tm_ed = time.time()
        if self.lbl_message.IsShown():
            self.lbl_message.Hide()
        cpu_num = os.cpu_count()
        if cpu_num is None:
            thread_num = int((tm_ed - tm_st))
        else:
            thread_num = min(cpu_num, int((tm_ed - tm_st)))
        thread_num = max(thread_num, 1)
        self.pool_capture = ThreadPoolExecutor(thread_num)
        time.sleep(0.5)
        threading.Thread(target=self.start_thread_capture, args=[1, thread_num]).start()
        print("Thread Num:", thread_num)
        pass

    def continuous_capture(self):
        while self.flag_capture.is_set():
            self.capture_and_show()
        pass

    def start_thread_capture(self, interval=0.5, n=1):
        for _ in range(n):
            if self.pool_capture._shutdown:
                break
            self.pool_capture.submit(self.continuous_capture)
            time.sleep(interval)
        pass

    def stop_thread_capture(self):
        if isinstance(self.pool_capture, ThreadPoolExecutor):
            self.pool_capture.shutdown(True)
            self.pool_capture = None
        self.Destroy()
        pass


class DlgPictureWithAction(DlgPicture, RecordActionForDlgPicture):

    def __init__(self, parent, cv_img: np.ndarray, scale: float = None):
        DlgPicture.__init__(self, parent, cv_img=cv_img, scale=scale)
        RecordActionForDlgPicture.__init__(self)
        self.Bind(wx.EVT_MOTION, self.on_move)
        self.Bind(wx.EVT_LEFT_UP, self.on_mouse_up)
        self.Bind(wx.EVT_LEFT_DOWN, self.on_mouse_down)

    def set_points(self, pts, index=-1):
        if index >= len(self.device_point):
            return False
        if index < 0 and len(pts) == len(self.device_point):
            self.device_point = [(int(p[0]), int(p[1])) for p in pts]
            self.clicked_point = [(int(p[0] * self.window_scale), int(p[1] * self.window_scale)) for p in pts]
        else:
            self.device_point[index] = (int(pts[0]), int(pts[1]))
            self.clicked_point[index] = (int(pts[0] * self.window_scale), int(pts[1] * self.window_scale))
        pass

    def get_points(self, index=-1):
        if index >= len(self.device_point):
            return (0, 0)
        if index < 0:
            return self.device_point
        else:
            return self.device_point[index]
        pass

    def get_press_time(self):
        return self.press_time

    def set_press_time(self, time):
        self.press_time = float(time)


class DlgActionsSimulator(DlgScreen, RecordActionForDlgPicture):

    def __init__(self, parent, adb: AdbInterface, scale: float = None, size: tuple = None):
        DlgScreen.__init__(self, parent, adb, scale, size)
        RecordActionForDlgPicture.__init__(self)
        self.mouse_pos = (-1, -1)
        self.actions = []
        # event
        if self.adb is not None:
            self.lbl_message.Bind(wx.EVT_MOTION, self.on_move)
            self.lbl_message.Bind(wx.EVT_LEFT_UP, self.on_mouse_up)
            self.lbl_message.Bind(wx.EVT_LEFT_DOWN, self.on_mouse_down)
            self.Bind(wx.EVT_MOTION, self.on_move)
            self.Bind(wx.EVT_LEFT_UP, self.on_mouse_up)
            self.Bind(wx.EVT_LEFT_DOWN, self.on_mouse_down)

    def set_image(self, cv_img: np.ndarray):
        if cv_img is not None:
            self.window_size = self.GetClientSize().GetWidth(), self.GetClientSize().GetHeight()
            self.cv_img_original = cv_img
            self.window_scale = self.window_size[0] / cv_img.shape[1]
            self.image_scale = 1. / self.window_scale
            self.cv_img_draw = cv2.resize(cv_img, self.window_size)
        else:
            self.cv_img_draw = self.cv_img_original = None
        self.repaint_mark(self.mouse_pos)
        pass

    def on_mouse_up(self, e: wx.MouseEvent):
        if RecordActionForDlgPicture.on_mouse_up(self, e):
            if self.device_point[0] == self.device_point[1]:
                self.actions.append([self.device_point[0], None, round(self.press_time, 3)])
                if self.press_time < 0.5:
                    self.adb.click(*self.device_point[0])
                else:
                    self.adb.press(*self.device_point[0], self.press_time)
            else:
                self.actions.append([self.device_point[0], self.device_point[1], round(self.press_time, 3)])
                self.adb.slide(*self.device_point[0], *self.device_point[1], self.press_time)
            return True
        return False

    def on_mouse_down(self, e: wx.MouseEvent):
        if self.press_stop_time > 0:
            self.actions.append([None, None, round(time.time() - self.press_stop_time, 3)])
        return RecordActionForDlgPicture.on_mouse_down(self, e)

    def on_move(self, e: wx.MouseEvent):
        self.mouse_pos = (e.GetX(), e.GetY())
        return RecordActionForDlgPicture.on_move(self, e)

if __name__ == "__main__":
    # print("非入口")
    # exit(0)
    from core.adb_interface import AdbInterface

    app = wx.App()
    adb = AdbInterface()
    adb.set_device("SJQNW19409003155", "HW")
    dlg = DlgActionsSimulator(None, adb)
    dlg.ShowModal()
    app.ExitMainLoop()
